import os
import subprocess
from collections import defaultdict


class SourcelineAnalyzer:
    def __init__(self, llvm_bin_path, so_files_path):
        self.llvm_bin_path = llvm_bin_path
        self.so_files = so_files_path
        self.dump_source = {}
        self.dump_percentage = defaultdict(lambda: defaultdict(float))
        self.addr_to_lineinfo = defaultdict(lambda: defaultdict(int))
        self.addr_to_event = defaultdict(lambda: defaultdict(int))

    def get_so_path(self, so_file):
        so_file = so_file.strip().split('/')[-1]
        for files in os.listdir(self.so_files):
            if so_file in files:
                return os.path.join(self.so_files, files)

    def add_event(self, dump_file_path, dump_addr, source_line, event_count, event_percentage):
        if dump_file_path not in self.dump_source:
            self.get_dumpfile(dump_file_path)
        dump_addr = dump_addr[2:]
        self.dump_percentage[dump_file_path][dump_addr] += event_percentage
        self.addr_to_lineinfo[dump_file_path][dump_addr] = source_line
        self.addr_to_event[dump_file_path][dump_addr] += event_count

    def analyze(self, source_info):
        for line in source_info:
            source_line = line['source_file_path'] + ':' + line['source_line']
            self.add_event(line['dump_file_path'], line['dump_addr'], source_line, line['event_count'], line['event_percentage'])
        return self.line_assembly()

    def get_dumpfile(self, dump_file_path):
        so_file = self.get_so_path(dump_file_path)
        cmd = [os.path.join(self.llvm_bin_path, 'llvm-objdump.exe'), '-d', so_file]
        result = subprocess.run(cmd, capture_output=True, text=True).stdout.strip()
        result = result.split('\n\n')[1:]
        result = [item for item in result if 'Disassembly of section' not in item]
        self.dump_source.update({dump_file_path: []})
        for line in result:
            self.dump_source.get(dump_file_path).append(line)

    def process_block_info(self, block, dump_file_path):
        info_processed = ''
        source_info = set()
        for line in block.split('\n'):
            addr = line.split(':')[0].lstrip(' ') 
            if addr not in self.addr_to_lineinfo[dump_file_path]:
                if line[-2:] == '>:':
                    info_processed += line + '\n'
                else:
                    info_processed += '\t\t\t' + line + '\n'
            else:
                info_processed += (self.addr_to_lineinfo[dump_file_path][addr].split(':')[-1] +
                                   '\t' + str(self.addr_to_event[dump_file_path][addr]) +
                                   '\t' + f"{self.dump_percentage[dump_file_path][addr]:.2f}%" +
                                   '\t' + line + '\n')
                source_info.add(str(self.addr_to_lineinfo[dump_file_path][addr]))
        
        return info_processed, source_info

    def line_assembly(self):
        ret = defaultdict(str)
        for dump_file_path, file_source in self.dump_source.items():
            for block in file_source:
                info_processed, source_info = self.process_block_info(block, dump_file_path)
                for source_file in source_info:
                    ret[source_file] += info_processed + '\n'
        return ret




